"""
calculate model size and FLOPs
"""
import sys
sys.path.append('./')

from easydict import EasyDict as edict
import torch

from mmcv import Config
from pcnets.models import build_model
from ptflops import get_model_complexity_info



def input_fn_partseg(*args, **kwargs):
    return dict(
        data = {'points': torch.randn([1, 3, 2048]), 'cls_label': torch.zeros([1], dtype=torch.int64)}
    )

def input_fn_classification(*args, **kwargs):
    return dict(
        data = {'points': torch.randn([1, 3, 1024])}
    )

args = edict()
# args.config = 'configs/ours/interpcnn2.py'
# args.config = 'configs/pointnet2/pointnet2ssg_cls_model40ply.py'
args.config = 'configs/ours/interpcnn2_partseg.py'
# args.config = 'configs/pointnet2/pointnet2msg_partseg.py'
args.inp_fn = input_fn_partseg

def main():
    cfg = Config.fromfile(args.config)

    model = build_model(cfg.model)
    model.init_weights()

    print(get_model_complexity_info(model, (None,), input_constructor=args.inp_fn))
    

if __name__ == '__main__':
    main()
